# check_data.py
import pickle
import pandas as pd


def check_data_files():
    # 检查 stock_data.pkl
    try:
        with open('stock_data.pkl', 'rb') as f:
            stock_data = pickle.load(f)
        print(f"stock_data.pkl 内容类型: {type(stock_data)}")
        print(f"包含的股票数量: {len(stock_data)}")

        if len(stock_data) > 0:
            print("前5个股票代码:", list(stock_data.keys())[:5])
            # 检查第一个股票的数据
            first_stock = list(stock_data.keys())[0]
            print(f"股票 {first_stock} 的数据形状: {stock_data[first_stock].shape}")
            print(f"股票 {first_stock} 的列名: {stock_data[first_stock].columns.tolist()}")
            print(
                f"股票 {first_stock} 的日期范围: {stock_data[first_stock].index.min()} 到 {stock_data[first_stock].index.max()}")
    except Exception as e:
        print(f"检查 stock_data.pkl 时出错: {e}")

    # 检查 trading_signals.pkl
    try:
        with open('trading_signals.pkl', 'rb') as f:
            trading_signals = pickle.load(f)
        print(f"\ntrading_signals.pkl 内容类型: {type(trading_signals)}")
        print(f"信号数量: {len(trading_signals)}")

        if len(trading_signals) > 0:
            print("前5个信号日期:", list(trading_signals.keys())[:5])
            first_date = list(trading_signals.keys())[0]
            print(f"日期 {first_date} 的信号: {trading_signals[first_date]}")
    except Exception as e:
        print(f"检查 trading_signals.pkl 时出错: {e}")


if __name__ == "__main__":
    check_data_files()